import torch
import torch.nn.functional as F
import math

def attention(query, key, value, mask=None, dropout=None):
    """计算'缩放点积注意力'"""

    # 首先取query的最后一维的大小，对应词嵌入维度
    d_k = query.size(-1)

    # 按照注意力公式，将query与key的转置相乘，这里面key将最后两个维度进行转置，再除以缩放系数得到注意力得分张量scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # 接着判断是否使用掩码张量
    if mask is not None:
        # 使用tensor的masked_fill方法，将掩码张量和scores张量每个位置一一比较，如果掩码张量则对应的scores张量用-1e9这个异常来替换
        scores = scores.masked_fill(mask == 0, -1e9)

    # 对scores的最后一维进行softmax操作，使用F.softmax方法，这样来得最终的注意力张量
    p_attn = F.softmax(scores, dim=-1)

    # 之后判断是否使用dropout进行随机置0
    if dropout is not None:
        p_attn = dropout(p_attn)

    # 最后，根据公式p_attn与value张量相乘来得最终的query注意力表示，同时返回注意力张量
    return torch.matmul(p_attn, value), p_attn


if __name__ == "__main__":
    # 测试attention函数
    print("=== 缩放点积注意力机制测试 ===")

    # 设置测试参数
    batch_size = 2
    seq_len = 5
    d_k = 64
    h = 8  # 头数

    # 创建测试数据
    query = torch.randn(batch_size, h, seq_len, d_k)
    key = torch.randn(batch_size, h, seq_len, d_k)
    value = torch.randn(batch_size, h, seq_len, d_k)

    print(f"输入张量shapes:")
    print(f"  query: {query.shape}")  # [2, 8, 5, 64]
    print(f"  key: {key.shape}")      # [2, 8, 5, 64]
    print(f"  value: {value.shape}")  # [2, 8, 5, 64]

    # 测试无mask的情况
    print(f"\n--- 测试1: 无mask ---")
    output, attn_weights = attention(query, key, value)
    print(f"输出output shape: {output.shape}")        # [2, 8, 5, 64]
    print(f"注意力权重 shape: {attn_weights.shape}")   # [2, 8, 5, 5]
    print(f"注意力权重和 (应该为1): {attn_weights.sum(dim=-1)[0, 0, 0].item():.4f}")

    # 测试有mask的情况
    print(f"\n--- 测试2: 有mask ---")
    # 创建mask：最后一个位置为0（被遮盖）
    mask = torch.ones(batch_size, 1, 1, seq_len)
    mask[:, :, :, -1] = 0  # 遮盖最后一个位置
    print(f"mask shape: {mask.shape}")  # [2, 1, 1, 5]
    print(f"mask内容: {mask[0, 0, 0]}")  # [1, 1, 1, 1, 0]

    output_masked, attn_weights_masked = attention(query, key, value, mask=mask)
    print(f"有mask的输出 shape: {output_masked.shape}")
    print(f"有mask的注意力权重 shape: {attn_weights_masked.shape}")
    print(f"最后一个位置的注意力权重 (应该接近0): {attn_weights_masked[0, 0, 0, -1].item():.6f}")
    print(f"注意力权重和 (应该为1): {attn_weights_masked.sum(dim=-1)[0, 0, 0].item():.4f}")

    # 测试有dropout的情况
    print(f"\n--- 测试3: 有dropout ---")
    dropout_layer = torch.nn.Dropout(p=0.1)
    output_dropout, attn_weights_dropout = attention(query, key, value, dropout=dropout_layer)
    print(f"有dropout的输出 shape: {output_dropout.shape}")
    print(f"dropout不改变注意力权重和: {attn_weights_dropout.sum(dim=-1)[0, 0, 0].item():.4f}")

    print(f"\n=== 测试完成 ===")